package LibDL.eval.regression;

import LibDL.core.Scalar;
import LibDL.core.Tensor;
import LibDL.eval.SupervisedEvaluator;

public abstract class AbstractRegressionEvaluator implements SupervisedEvaluator {
    private Tensor truth;
    private Tensor pred;
    private Tensor sampleWeight;
    private Tensor outputWeight;

    @Override
    public double evaluate(Tensor truth,Tensor pred){
        setTruth(truth);
        setPred(pred);
        return evaluate();
    }

    @Override
    public double evaluate() {
        return core(getTruth(), getPred());
    }

    abstract protected double core(Tensor truth,Tensor pred);

    /*
     ********** tool methods **********
     */
    protected enum TargetType {
        SINGLE_OUTPUT, MULTI_OUTPUT, NONE
    }
    protected TargetType checkTarget(Tensor truth, Tensor pred){
        if (truth.sizes().size() == 1 && pred.sizes().size() == 1
                && truth.size(0) == pred.size(0)){
            return TargetType.SINGLE_OUTPUT;
        } else if (truth.sizes().size() == 2 && pred.sizes().size() == 2
                && truth.size(0) == pred.size(0)
                && truth.size(1) == pred.size(1)){
            return TargetType.MULTI_OUTPUT;
        }
        return TargetType.NONE;
    }

    protected void checkConsistentShape(Tensor truth,Tensor sampleWeight,
                                             Tensor outputWeight,
                                             TargetType type){
        if (truth.size(0) != sampleWeight.size(0)){
            throw new IllegalArgumentException("The length of sample weight's first dim" +
                    " must be consistent with that of truth and pred." +
                    " Expect: "+truth.size(0)+". Found: "+sampleWeight.size(0));
        }else if (type.equals(TargetType.MULTI_OUTPUT)
                && truth.size(1) != outputWeight.size(0)){
            throw new IllegalArgumentException("The length of output weight must be consistent" +
                    " with the second dim length of truth and pred." +
                    " Expected: "+truth.size(1)+" Found: "+outputWeight.size(0));
        }
    }

    /**
     * @param sampleScore 1-dim tensor
     * @param sampleWeight 1-dim tensor
     * @param normalize whether return weighted sum or avg of scores
     * @return double
     */
    protected Tensor weightedSum(Tensor sampleScore, Tensor sampleWeight, boolean normalize, int axis){
        if (normalize){
            sampleWeight = sampleWeight.div(new Scalar(sampleWeight.sum().item().to_double()));
        }
        if (sampleScore.sizes().size() == 1) {
            return sampleScore.mul(sampleWeight).sum(axis);
        } else {
            return sampleScore.mul(sampleWeight.unsqueeze(1-axis)).sum(axis);
        }
    }

    /*
     ********** getter & setter **********
     */
    public Tensor getTruth() {
        return truth;
    }

    public AbstractRegressionEvaluator setTruth(Tensor truth) {
        this.truth = truth;
        return this;
    }

    public Tensor getPred() {
        return pred;
    }

    public AbstractRegressionEvaluator setPred(Tensor pred) {
        this.pred = pred;
        return this;
    }

    public Tensor getSampleWeight() {
        return sampleWeight;
    }

    public AbstractRegressionEvaluator setSampleWeight(Tensor sampleWeight) {
        this.sampleWeight = sampleWeight;
        return this;
    }

    public Tensor getOutputWeight() {
        return outputWeight;
    }

    public AbstractRegressionEvaluator setOutputWeight(Tensor outputWeight) {
        this.outputWeight = outputWeight;
        return this;
    }
}
